from tqdm import tqdm
import torch
from matplotlib import pyplot as plt
import numpy as np
def save_image(image, filename):
    # 确保图像在 CPU 上
    image = image.to('cpu')
    # 转换 Tensor 为 numpy 数组并调整维度为 (H, W, C)
    image = image.detach().numpy().transpose(1, 2, 0)
    # 如果图像数据在 [-1, 1] 范围内，将其转换到 [0, 1]
    if image.min() < 0:
        image = (image + 1) / 2
    # 裁剪数值以确保它们在 [0, 1] 范围内
    image = np.clip(image, 0, 1)
    # 保存图像
    plt.imsave(filename, image)



def train(net, optimizer, train_loader, loss_fc, epoch):
    net.train()
    device='cuda' if torch.cuda.is_available() else 'cpu'
    net.to(device)
    for e in range(epoch):
        print(f"Training epoch {e}")
        error=0.0
        for d, data in enumerate(tqdm(train_loader)):
            optimizer.zero_grad()
            data=data.to(device)
            output, mean, var=net(data)
            loss=loss_fc(data, output, mean, var)
            loss.backward()
            optimizer.step()
            error+=loss.item()
            if d%400==0 and e==1:
                save_image(image=output[0], filename=f'{d}.jpg')
        error/=len(train_loader)
        print(f"Epoch {e}, Error: {error:.4f}")
        if e==250:
            torch.save(net.state_dict(), '250.pth')  # 同上
    print("Training completed")
    torch.save(net.state_dict(), 'net_params.pth')  # 同上
    print('param saved')
    
def test(net, test_loader, loss_fc):
    net.eval()
    device='cuda' if torch.cuda.is_available() else 'cpu'
    net.to(device)
    error=0.0
    for d, data in enumerate(test_loader):
        data.to(device)
        output, mean, var=net(data)
        loss=loss_fc(data, output, mean, var)
        error+=loss.item()
    error.to('cpu')
    error/=len(test_loader)
    print(f"Test Loss: {error}")